import json

from dashscope import Generation
from env import key

shelf_mapper = {
    "A": (3.0, 4.0),
    "B": (3.0, -3.8),
    "JD": (-4.0, 4.0),
}

# 预定义示例响应（用于few-shot提示）
example1_response = json.dumps(
    {
        "x": 3.0,
        "y": 4.0
    },
    ensure_ascii=False
)
example2_response = json.dumps(
    {
        "x": 3.0,
        "y": -3.8
    },
    ensure_ascii=False
)
example3_response = json.dumps(
    {
        "x": -4.0,
        "y": -4.0
    },
    ensure_ascii=False
)


messages = [
    {
        "role": "system",
        "content": f"""
        #背景#
        你作为快递站的决策系统，要根据{shelf_mapper}这一映射关系，将输入的快递取件码，映射到对应快递柜的坐标，并输出包含x和y的JSON。
        #目的#
        输出映射后快递柜的坐标
        #任务步骤#
        1. 从输入中提取快递取件码。
        2. 从快递取件码中提取快递柜ID。
        3. 将快递柜ID作为键，在{shelf_mapper}中找到对应的坐标值，如果没找到使用默认值"x": 0.0, "y": 0.0。
        4. 将坐标元组(x, y)转化为目标JSON格式输出。
        #示例#
        Q: 放到A货架。
        A: {example1_response}
        Q: 把5B-46839送到货架。
        A: {example2_response}
        Q: 把货送到JD货架。
        A: {example3_response}
        """
    }
]


def task(question):
    print(f"Sending question: {question}")
    messages.append({
        "role": "user",
        "content": question
    })
    response = Generation.call(
        api_key=key,
        model="qwen-plus",  # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
        messages=messages,
        result_format='message',
        response_format={'type': 'json_object'}
    )
    json_string = response.output.choices[0].message.content
    messages.remove({
        "role": "user",
        "content": question
    })
    print(json_string)

    data = json.loads(json_string)
    x = float(data['x'])
    y = float(data['y'])
    return x, y
